import os
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from main import CONFIG

db_config = CONFIG['database']

if db_config['type'] == 'sqlite':
    engine = create_engine(f"sqlite:///{db_config['dbname']}")

elif db_config['type'] == 'mysql':
    engine = create_engine(
        f"mysql+pymysql://{db_config['username']}:{db_config['password']}"
        f"@{db_config['host']}:{db_config['port']}/{db_config['dbname']}"
    )

elif db_config['type'] == 'postgresql':
    engine = create_engine(
        f"postgresql+psycopg2://{db_config['username']}:{db_config['password']}"
        f"@{db_config['host']}:{db_config['port']}/{db_config['dbname']}"
    )

SessionLocal = sessionmaker(bind=engine)

def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()